# 通过connect_unity得到的unity_env比起一般的强化学习环境(比如openAI的gym)来说还是复杂了一些。
# 所以这里选择基于unity_env再定义一些基本的模板（就是把unity_env的各种情况针对性地分解开），使用起来会简单些。
import collections
import copy
import random
import numpy as np
from tools import connect_unity as cu


class ReplayBuffer:
    """ 经验回放池 """
    def __init__(self, capacity):
        self.buffer = collections.deque(maxlen=capacity)  # 队列,先进先出

    def add(self, state, action, reward, next_state, done):  # 将数据加入buffer
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):  # 从buffer中采样数据,数量为batch_size
        transitions = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = zip(*transitions)
        return np.array(state), np.array(action), reward, np.array(next_state), done

    def size(self):  # 目前buffer中数据的数量
        return len(self.buffer)

    def clear(self):  # 清空buffer中的所有数据
        self.buffer.clear()


class Cell(object):
    """每个Cell最多存储一条（state,action,reward,next_state）"""

    def __init__(self, cell_len=4):
        self.cell_len = cell_len
        self.store = []
        self.current_sym = None
        self.types = ['state', 'action', 'reward']
        self.skip_next = False  # 如果不是调用reset重置环境，它自己重置会返回 reward（无意义）和 state,那么就要跳过这个无意义的reward

        # 用来记录每幕的总回报的部分
        self.ep_return = 0
        self.all_returns = []

    def next_sym(self):
        if self.current_sym is None or self.current_sym == self.types[2]:
            return self.types[0]
        elif self.current_sym == self.types[0]:
            return self.types[1]
        elif self.current_sym == self.types[1]:
            return self.types[2]

    def add_data(self, data, sym, ep_end):
        """
        :param ep_end: 标记是否为幕的最后一个状态。
        :param sym: 标记了当前数据是什么 ['state','action','reward']
        :param data: 必须是state, action, reward 其中之一
        """
        if self.skip_next:
            self.skip_next = False
            if sym != self.types[2]:
                raise ValueError(f"选择跳过该条数据，但是该条数据不是{self.types[2]}，而是{sym}")
            return None

        if len(self.store) >= self.cell_len:
            raise ValueError("Cell已满，无法继续添加！")
        if sym != self.next_sym():
            print(self.store)
            raise ValueError(f"期望类型{self.next_sym()},但是得到了类型{sym}")
        self.store.append(data)
        self.current_sym = sym

        if sym == self.types[2]:  # 要记录该奖励。
            self.ep_return += data

        r_data = None
        if sym == self.types[0]:  # 表示当前输入的数据是状态，可以进一步判断是否可以收集数据
            if len(self.store) >= self.cell_len:
                r_data = copy.deepcopy(self.store)
                self.store = self.store[-1:]  # 保留 next_state
                self.current_sym = self.types[0]
        if ep_end:
            self.store = []
            self.current_sym = None
            self.skip_next = True
            self.all_returns.append(self.ep_return)
            self.ep_return = 0

        return r_data

    def clear_cell(self):
        self.store = []
        self.current_sym = None
        self.skip_next = False
        self.ep_return = 0
        self.all_returns = []


class EnvironmentV1(object):
    """
    V1版本说明：
    observation是一个向量，action也是一个向量，且action是连续的。
    只考虑一个behavior_name (环境当中有多个就选择第一个)。
    """
    def __init__(self, time_scale, no_graphics, agent_num, buffer_size, exe_path=None, port=5004, timeout=30):
        """
        :param exe_path: 可执行文件路径（None表示编辑器模式下训练）
        :param port: 通信端口（默认5004）
        :param timeout: 最大等待时间(秒)
        :param time_scale: 控制物理模拟的时间流速
        :param no_graphics: 如果为True,那么使用打包好之后的程序进行训练就不会出现图形界面
        """
        self.unity_env = cu.connect_unity(exe_path=exe_path, port=port, timeout=timeout, time_scale=time_scale,
                                     no_graphics=no_graphics)
        self.behavior_name = None
        self.observation_dim = None
        self.action_dim = None
        self.max_action = None
        self.exp_buffer = ReplayBuffer(buffer_size)
        self.agent_num = agent_num
        self.recorder = [Cell() for _ in range(agent_num)]

        self.initialize()

    def initialize(self):
        unity_env_infor = cu.get_unity_env_info(self.unity_env)
        unity_env_infor = unity_env_infor[0] # 取第一个behavior_name对应的信息。

        self.behavior_name = unity_env_infor["behavior_name"]
        self.observation_dim = unity_env_infor["observation_dim"][0][0]
        self.action_dim = unity_env_infor["action_dim"][0]
        self.max_action = 1 #这个参数设置为1，但是还可以在Unity端进行缩放。

    def add_exp(self, datas, agent_ids, sym, ep_end):
        """
        datas 可以是 state,action,reward 其中之一
        agent_ids 指定了 datas 对应的agent是哪些
        sym 表示当前加入的数据具体是 state,action,reward 当中的哪一个
        ep_end 用来表示 datas 是否为 幕的最后一个状态(时间限制)。
        """
        for idx, agent_id in enumerate(agent_ids):
            """ 
            idx 是 agent_id 对应的数据在datas当中的索引
            agent_id 表示了 这条数据要插入 self.recorder 的几号位置。
            """
            exp = self.recorder[agent_id].add_data(data=datas[idx],sym=sym,ep_end=ep_end)
            if exp is not None:
                self.exp_buffer.add(state=exp[0], action=exp[1], reward=exp[2], next_state=exp[3], done=ep_end)

    def reset(self):
        """
        重置Unity端场景当中的全部agent,然后返回self.behavior_name指定的agent的初始状态。
        :return: states (num_agent, state_dim), agent_ids (num_agents,)
        """
        states_ids = cu.reset_unity_env(self.unity_env,[self.behavior_name])
        states_ids = states_ids[0]
        states = states_ids["observations"][0] #(num_agent, state_dim)
        ids = states_ids["agent_ids"] #(num_agents,)

        for cell in self.recorder:
            cell.clear_cell() #环境重置了，那么记录器也要清空。

        self.add_exp(datas=states,agent_ids=ids,sym=self.recorder[0].types[0],ep_end=False)

        return states, ids

    def step(self, continuous_actions, agent_ids):
        """
        :param continuous_actions: 具体需要发送的动作 (num_agent, action_dim)
        :param agent_ids: 这些动作要发送给哪些 id 对应的 agent
        """
        alive, dead = cu.step_unity_env(self.unity_env, self.behavior_name, agent_ids, continuous_actions)
        # 会出现某个agent同时出现在alive和dead的情况（一般出现在某个幕触发EndEpisode()的时候），优先认为该agent属于alive部分。
        alive_obs, alive_rewards, alive_ids = alive
        dead_obs, dead_rewards, dead_ids = dead
        if len(alive_ids)+len(dead_ids)>self.agent_num: #确定出现了重复。
            duplicate_mask = np.isin(dead_ids, alive_ids)
            dead_obs = [obs[~duplicate_mask] for obs in dead_obs] if len(dead_obs) > 0 else []
            dead_rewards = dead_rewards[~duplicate_mask]
            dead_ids = dead_ids[~duplicate_mask]

        # ===========下面开始收集经验===============
        self.add_exp(datas=continuous_actions,agent_ids=agent_ids,sym=self.recorder[0].types[1],ep_end=False) # actions

        self.add_exp(datas=alive_rewards,agent_ids=alive_ids,sym=self.recorder[0].types[2],ep_end=False) # rewards
        self.add_exp(datas=dead_rewards, agent_ids=dead_ids,sym=self.recorder[0].types[2],ep_end=False) # rewards

        self.add_exp(datas=alive_obs[0],agent_ids=alive_ids,sym=self.recorder[0].types[0],ep_end=False) #next_state
        self.add_exp(datas=dead_obs[0], agent_ids=dead_ids,sym=self.recorder[0].types[0],ep_end=True) #next_state

        next_state = alive_obs[0]
        next_ids = alive_ids
        return next_state, next_ids



